// 

#ifndef SIMPLE_MODULES_SIMPLEGPU_H_
#define SIMPLE_MODULES_SIMPLEGPU_H_

#include <queue>
#include <utility>
#include <unordered_map>
#include <vector>
#include "cpu.h"
#include "ib_cluster/base/ai_sim.h"

namespace ai_sim {
class Cpu;

class Gpu : public omnetpp::cSimpleModule {
 private:
    int my_index;
    int my_global_index;
    int numNvlink_port;
    int numComm_peer;
    std::unordered_map<int, std::queue<omnetpp::cMessage*>> msg_queues;
    Cpu *my_cpu;

 public:
    void initialize() override;
    void handleMessage(omnetpp::cMessage *msg) override;
    void handleSend(omnetpp::cMessage *msg);
    void handleSend(Flow);
};
}  // namespace ai_sim

#endif  // SIMPLE_MODULES_SIMPLEGPU_H_
